import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D

# 加载深度图（假设深度图是 [H, W] 的 .npy 文件）
depth_array = np.load("norm_test/gt_depth.npy")  # 替换为你的深度图路径
depth_array2 = np.load("norm_test/pred_depth.npy")  # 替换为你的深度图路径
# depth_array3 = np.load("norm_test/recovered_depth.npy")  # 替换为你的深度图路径

# 获取深度图的尺寸
height, width = depth_array.shape

# 创建网格坐标
x = np.arange(width)
y = np.arange(height)
x, y = np.meshgrid(x, y)


# 将深度值归一化（假设深度值范围是 0-最大值）
z = depth_array / np.max(depth_array)
z2 = depth_array2 / np.max(depth_array2)
# z3 = depth_array3 / np.max(depth_array3)

# 创建 3D 图形
fig = plt.figure()
ax = fig.add_subplot(111, projection="3d")

# 绘制点云
ax.scatter(x, y, z, c=z, cmap="viridis", s=1)  # 使用深度值作为颜色映射
ax.scatter(x, y, z2, c=z2, cmap="spring",alpha=0.2, s=1)  # 使用深度值作为颜色映射
# ax.scatter(x, y, z3, c=z3, cmap="winter",alpha=0.2, s=1)  # 使用深度值作为颜色映射

# 设置坐标轴标签
ax.set_xlabel("X")
ax.set_ylabel("Y")
ax.set_zlabel("Z")

# 设置视角（可选）
ax.view_init(elev=20, azim=35)  # 调整视角

# 保存为图片
plt.savefig("depth_point_cloud.png", dpi=300)  # dpi 控制图片分辨率

# 显示图形（可选）
plt.show()

print("3D 点云图片已保存为 depth_point_cloud.png")